#pragma once
#include <map>
#include <string>
#include <tuple>
#include <memory>
#include <d3d11.h>
#include "CComPtrEx.h"
#include "Resources.h"
#include "ShaderCache.h"

#define SHADER_VERSION _T("5")

template<typename ShaderInterface> struct ShaderProfile { static const TCHAR * name; };

void InitShaderProfiles(D3D_FEATURE_LEVEL fl);

class ShaderManager
{
private:
	CComPtrEx<ID3D11Device> m_device;
	bool m_keepAllBytecode;
	std::unique_ptr<ShaderCache> m_shaderCache;

	class IElement
	{
	public:
		IElement() {}
		virtual ~IElement() {}
	};
	template<typename T> class Element : public IElement
	{
	public:
		Element() { }
		virtual ~Element() { }
		CComPtrEx<T> element;
		CComPtrEx<ID3DBlob> shaderCode;
	};

	std::map<std::string, IElement*> m_storage;

	template<typename T>
	bool CreateShaderObject(const std::string& key, CComPtrEx<ID3DBlob>& data, Element<T>* result)
	{
		return false;
	}
	template<> bool CreateShaderObject(const std::string& key, CComPtrEx<ID3DBlob>& data, Element<ID3D11VertexShader>* result)
	{
		auto res = SUCCEEDED(m_device->CreateVertexShader(data->GetBufferPointer(),
			data->GetBufferSize(), nullptr, &result->element));
		result->shaderCode = data;
		return res;
	}
	template<> bool CreateShaderObject(const std::string& key, CComPtrEx<ID3DBlob>& data, Element<ID3D11PixelShader>* result)
	{
		auto res = SUCCEEDED(m_device->CreatePixelShader(data->GetBufferPointer(),
			data->GetBufferSize(), nullptr, &result->element));
		if (m_keepAllBytecode)
		{
			result->shaderCode = data;
		}
		return res;
	}
	template<> bool CreateShaderObject(const std::string& key, CComPtrEx<ID3DBlob>& data, Element<ID3D11GeometryShader>* result)
	{
		auto res = SUCCEEDED(m_device->CreateGeometryShader(data->GetBufferPointer(),
			data->GetBufferSize(), nullptr, &result->element));
		if (m_keepAllBytecode)
		{
			result->shaderCode = data;
		}
		return res;
	}
	template<> bool CreateShaderObject(const std::string& key, CComPtrEx<ID3DBlob>& data, Element<ID3D11HullShader>* result)
	{
		auto res = SUCCEEDED(m_device->CreateHullShader(data->GetBufferPointer(),
			data->GetBufferSize(), nullptr, &result->element));
		if (m_keepAllBytecode)
		{
			result->shaderCode = data;
		}
		return res;
	}
	template<> bool CreateShaderObject(const std::string& key, CComPtrEx<ID3DBlob>& data, Element<ID3D11DomainShader>* result)
	{
		auto res = SUCCEEDED(m_device->CreateDomainShader(data->GetBufferPointer(),
			data->GetBufferSize(), nullptr, &result->element));
		if (m_keepAllBytecode)
		{
			result->shaderCode = data;
		}
		return res;
	}
	template<> bool CreateShaderObject(const std::string& key, CComPtrEx<ID3DBlob>& data, Element<ID3D11ComputeShader>* result)
	{
		auto res = SUCCEEDED(m_device->CreateComputeShader(data->GetBufferPointer(),
			data->GetBufferSize(), nullptr, &result->element));
		if (m_keepAllBytecode)
		{
			result->shaderCode = data;
		}
		return res;
	}

	void CreateDefinesString(const D3D10_SHADER_MACRO* pDefines, std::string& result)
	{
		if (!pDefines)
		{
			result.clear();
			return;
		}
		const D3D10_SHADER_MACRO* curr = pDefines;
		while ((curr->Name) && (curr->Definition))
		{
			result += curr->Name;
			result += ":";
			result += curr->Definition;
			result += ";";
			curr++;
		}
	}
public:
	ShaderManager(ID3D11Device * device, bool keepAllBytecode = false)
	{
		m_device = device;
		m_keepAllBytecode = keepAllBytecode;
		TCHAR path[MAX_PATH + 1];
		GetTempPath(MAX_PATH, path);
		std::string tempFolderPath = path;
		tempFolderPath += _T("D3D11Client_ShadersCache\\");
		auto attrs = ::GetFileAttributes(tempFolderPath.c_str());
		if (attrs == INVALID_FILE_ATTRIBUTES)
		{
			//directory does not exist, create it
			::CreateDirectory(tempFolderPath.c_str(), nullptr);
		}
		m_shaderCache.reset(new ShaderCache(tempFolderPath));
	}
	~ShaderManager()
	{
		for(auto it = m_storage.begin(); it != m_storage.end(); it++)
		{
			delete (it->second);
		}
		m_storage.clear();
	}

	template <typename T>
	bool GetShader(const std::string& fileName, const std::string& entryPoint, CComPtrEx<T>& result, const D3D10_SHADER_MACRO* pDefines = nullptr)
	{
		std::string shaderProfile = ShaderProfile<T>::name;
		std::string key = fileName + "__EP__" + entryPoint + "__SP__" + shaderProfile;
		std::string definesStr;
		CreateDefinesString(pDefines, definesStr);
		key += "__D__" + definesStr;
		if (m_storage.count(key) > 0)
		{
			auto val = dynamic_cast<Element<T>*>(m_storage[key]);
			result = val->element;
			return true;
		}
		else
		{
			CComPtrEx<ID3DBlob> data;
			CacheEntry entry;
			entry.keyLength = (USHORT)key.length();
			entry.key = key.c_str();
			if (m_shaderCache->SearchForEntry(entry, fileName))
			{
				//shader found in cache, so just load it
				D3DCreateBlob(entry.bytecodeLength, &data);
				CopyMemory(data->GetBufferPointer(), entry.bytecode, entry.bytecodeLength);
			}
			else
			{
				//shader is no in cache
				if (CompileShaderFromFile(fileName.c_str(), entryPoint.c_str(), shaderProfile.c_str(), &data, pDefines))
				{
					entry.bytecodeLength = (UINT)data->GetBufferSize();
					entry.bytecode = data->GetBufferPointer();
					m_shaderCache->AddCacheEntry(entry, fileName);
				}
				else
				{
					return false;
				}
			}
			auto element = new Element<T>();
			if (CreateShaderObject(key, data, element))
			{
				m_storage.insert(std::make_pair(key, element));
				result = element->element;
				return true;
			}
			delete element;
		}
		return false;
	}
	template<typename T>
	bool GetShaderBytecode(const CComPtrEx<T>& shaderObject, CComPtrEx<ID3DBlob>& bytecode)
	{
		for(auto it = m_storage.begin(); it != m_storage.end(); it++)
		{
			auto element = dynamic_cast<Element<T>*>(it->second);
			if (element && element->element == shaderObject)
			{
				bytecode = element->shaderCode;
				return true;
			}
		}
		return false;
	}
};

#define SAFE_RELEASE(p) if ((p)) { (p)->Release(); (p) = nullptr; }

#define BEGIN_LAYOUT(name) D3D11_INPUT_ELEMENT_DESC name[] = {
//#define DEFINE_LAYOUT_FIRST_ELEMENT(semanticName,format) {semanticName, 0, format, 0, 0, D3D11_INPUT_PER_VERTEX_DATA, 0 },
#define DEFINE_LAYOUT_ELEMENT(semanticName,format,index) {semanticName, index, format, 0, D3D11_APPEND_ALIGNED_ELEMENT, D3D11_INPUT_PER_VERTEX_DATA, 0 },
#define END_LAYOUT };
#define LAYOUTS_COUNT(name) sizeof(name)/sizeof(name[0])
